import chromadb
from langchain.vectorstores import Chroma
from langchain.schema import Document
from typing import List, Optional
import yaml
import os


class VectorStoreManager:
    """向量数据库管理器"""

    def __init__(self, embeddings, config_path: str = "config.yaml"):
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)

        self.config = config['vector_store']
        self.embeddings = embeddings
        self.vector_store = None
        self._initialize_store()

    def _initialize_store(self):
        """初始化向量存储"""
        persist_dir = self.config['persist_directory']
        os.makedirs(persist_dir, exist_ok=True)

        self.vector_store = Chroma(
            collection_name=self.config['collection_name'],
            embedding_function=self.embeddings,
            persist_directory=persist_dir
        )

    def add_documents(self, documents: List[Document]):
        """添加文档到向量数据库"""
        self.vector_store.add_documents(documents)
        self.vector_store.persist()

    def similarity_search(self, query: str, k: int = 5) -> List[Document]:
        """相似性搜索"""
        return self.vector_store.similarity_search(query, k=k)

    def similarity_search_with_score(self, query: str, k: int = 5):
        """带分数的相似性搜索"""
        return self.vector_store.similarity_search_with_score(query, k=k)